Perf: faster get_largest_connected_component_mask (bincount + LUT gather)#8978
Conversation
Rank component sizes with a full-field bincount (zeroing background at index 0) instead of gathering every non-background voxel through lib.nonzero, and build the output mask with a boolean lookup-table gather instead of lib.isin over the whole label field. Both avoid large transient allocations; output is bit-identical. Covers the numpy and cupy/cucim paths unchanged. Signed-off-by: Soumya Snigdha Kundu <soumya_snigdha.kundu@kcl.ac.uk>
📝 WalkthroughWalkthroughThe change refactors get_largest_connected_component_mask in monai/transforms/utils.py to compute label voxel counts via bincount on flattened features, zero out the background count, rank labels by descending count via argsort, and build a boolean keep lookup table indexed directly into the feature map, replacing the prior nonzero-based extraction and isin masking approach. Estimated code review effort: 2 (Simple) | ~10 minutes Changes
Related issues: Addresses Related PRs: None noted. Suggested labels: performance, transforms Suggested reviewers: MONAI transforms maintainers Poem: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
monai/transforms/utils.py (1)
1227-1235: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winUse
argpartitioninstead of fullargsortfor top-k selection.
lib.argsort(counts)[::-1][:num_components]sorts the entire counts array just to pick the topnum_components. Since only the largest few are needed,argpartitionavoids the full O(n log n) sort and better matches this PR's own perf-optimization goal.♻️ Proposed refactor
- # argsort[::-1] gives labels of the largest components; keep the first n - features_to_keep = lib.argsort(counts)[::-1][:num_components] + # argpartition avoids a full sort when only the top-n labels are needed + features_to_keep = lib.argpartition(counts, -num_components)[-num_components:]🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@monai/transforms/utils.py` around lines 1227 - 1235, The top-k selection in the label-count ranking is doing a full sort via lib.argsort(counts)[::-1][:num_components], which is unnecessary for selecting only the largest components. Update the selection logic in the component-ranking code around the bincount/counts handling to use lib.argpartition for the top num_components labels, then keep the existing boolean lookup-table flow (keep/features_to_keep/out) to preserve behavior while improving performance.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@monai/transforms/utils.py`:
- Around line 1227-1235: The KeepLargestConnectedComponent ranking path in
utils.py is not covered for equal-sized connected components, so the current
argsort-based selection can change silently. Add a regression test for the
KeepLargestConnectedComponent behavior with tie-sized components that exercises
the ranking branch and asserts the chosen components remain deterministic; use
the same multi-component setup already used in the existing connected-component
tests to locate the relevant case.
---
Nitpick comments:
In `@monai/transforms/utils.py`:
- Around line 1227-1235: The top-k selection in the label-count ranking is doing
a full sort via lib.argsort(counts)[::-1][:num_components], which is unnecessary
for selecting only the largest components. Update the selection logic in the
component-ranking code around the bincount/counts handling to use
lib.argpartition for the top num_components labels, then keep the existing
boolean lookup-table flow (keep/features_to_keep/out) to preserve behavior while
improving performance.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 26c10d81-e6f4-4fe3-b7ad-e8038af614ce
📒 Files selected for processing (1)
monai/transforms/utils.py
|
Im choosing not to add that test because this code change does not change behaviour from the previous one. This whole tie angle i believe isnt something to really factor in. |
2 participants
Fixes #8974 .
Description
get_largest_connected_component_mask(monai/transforms/utils.py) ranked component sizes by gathering every non-background voxel throughlib.nonzero(features)and then built the mask withlib.isinover the whole label field. Both allocate large transient arrays. This counts labels with a single full-fieldlib.bincount(zeroing index 0 to drop background), and builds the mask with a boolean lookup-table gather (keep[features]) instead ofisin. Output is bit-identical, verified across 2D and 3D at several foreground fractions and component counts, and the numpy and cupy/cucim paths are covered unchanged.The
isinreplacement is the dominant win on both time and peak memory; thebincountchange removes the redundantnonzeroindex arrays and gathered copy on top.Types of changes